
#include "imageSegmentator.h"

#include <QtCore/QString>
#include <QtGui/QPainter>

/*----------------------------------------------------------------------------*/


ImageSegmentator::ImageSegmentator(		QImage& _image
									,	int _imageSize
									,	int _imageContrast
									,	float _lineExpFactor
									,	float _lineTopFactor
									,	float _lineBotFactor
									,	float _wordLeftFactor
									,	float _wordRightFactor
									,	float _symbIntervalFactor
									,	float _symbBrightFactor
									,	float _symbWidthFactor
									)
	:	m_image( _image )
	,	m_linesFrames()
	,	m_wordsFrames()
	,	m_symbolsFrames()
	,	m_linesImages()
	,	m_wordsImages()
	,	m_symbolsImages()		
	,	m_symbolsLayout()
{
	/** finding lines frames */
	findLinesFrames( _lineExpFactor, _lineTopFactor, _lineBotFactor );

	/** generate lines images */
	generateLinesImages();

	/** finding words frames */
	findWordsFrames( _imageContrast, _wordLeftFactor, _wordRightFactor );

	/** generate words images */
	generateWordsImages();

	/** finding symbol frames */
	findSymbolsFrames( _symbIntervalFactor, _symbBrightFactor, _symbWidthFactor );

	/** generate symbol images */
	generateSymbolsImages( _imageSize, _imageContrast );
}


/*----------------------------------------------------------------------------*/


ImageSegmentator::~ImageSegmentator()
{
}


/*----------------------------------------------------------------------------*/


QRectV& ImageSegmentator::GetLinesFrames()
{
	return m_linesFrames;
}


/*----------------------------------------------------------------------------*/


QRectV& ImageSegmentator::GetWordsFrames()
{
	return m_wordsFrames;
}


/*----------------------------------------------------------------------------*/


QRectV& ImageSegmentator::GetSymbolsFrames()
{
	return m_symbolsFrames;
}


/*----------------------------------------------------------------------------*/


QImageV& ImageSegmentator::GetSymbolsImages()
{
	return m_symbolsImages;
}


/*----------------------------------------------------------------------------*/


IntVV& ImageSegmentator::GetSymbolsLayout()
{
	return m_symbolsLayout;
}


/*----------------------------------------------------------------------------*/


float ImageSegmentator::getPixelGreyBright( QImage& _image, int _i, int _j )
{
	/** returns: 0 - white, 255 - black */
	return (float)( abs( qGray( _image.pixel( _j, _i ) ) - 255 ) );	
}


/*----------------------------------------------------------------------------*/


FloatV ImageSegmentator::getImageRowsBrightness( QImage& _image )
{
	FloatV rowsBrightness( _image.height(), 0.0 );
	
	for ( int i = 0; i < _image.height(); i++ )
	{
		for ( int j = 0; j < _image.width(); j++ )
		{
			rowsBrightness[i] += getPixelGreyBright( _image, i, j );
		}
		rowsBrightness[i] /= (float)_image.width();		
	}
	return rowsBrightness;
}


/*----------------------------------------------------------------------------*/


FloatV ImageSegmentator::getImageColsBrightness( QImage& _image )
{
	FloatV colsBrightness( _image.width(), 0.0 );

	for ( int i = 0; i < _image.width(); i++ )
	{
		for ( int j = 0; j < _image.height(); j++ )
		{
			colsBrightness[i] += getPixelGreyBright( _image, j, i );
		}
		colsBrightness[i] /= (float)_image.height();
	}
	return colsBrightness;
}


/*----------------------------------------------------------------------------*/


float ImageSegmentator::getAverageBrightness( FloatV& _brightness )
{
	float averageBrightness = 0.0;
	for ( size_t i = 0; i < _brightness.size(); i++ )
	{
		averageBrightness += _brightness[i];
	}
	averageBrightness /= (float)_brightness.size();	
	return averageBrightness;
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::setPixelGreyBright( QImage& _image, int _i, int _j, int _color )
{
	int color = abs( _color - 255 );
	_image.setPixel( _j, _i, qRgb( color, color, color ) );
}


/*----------------------------------------------------------------------------*/


QImage ImageSegmentator::contrastIncrease( QImage& _image, int _contrast )
{
	QImage tempImage = _image;

	for ( int i = 0; i < _image.height(); i++ )
	{			
		for ( int j = 0; j < _image.width(); j++ )
		{			
			if ( getPixelGreyBright( _image, i, j ) > _contrast )			
				setPixelGreyBright( tempImage, i, j, 255 );			
			else			
				setPixelGreyBright( tempImage, i, j, 0 );						
		}
	}

	return tempImage;
}


/*----------------------------------------------------------------------------*/


QImage ImageSegmentator::imageDegradation( QImage& _image )
{
	QImage tempImage = _image;
	for ( int i = 0; i < _image.height(); i++ )
	{
		for ( int j = 0; j < _image.width(); j++ )
		{
			if ( getPixelGreyBright( _image, i, j ) != 0 )
				paintNearbyPixels( tempImage, i, j );
		}
	}
	return tempImage;
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::paintNearbyPixels( QImage& _image, int _i, int _j )
{
	int border = 3;
	for ( int i = _i - border; i <= _i + border; i++ )
	{
		for ( int j = _j - border; j <= _j + border; j++ )
		{
			if( 0 <= i && i < _image.height() && 0 <= j && j < _image.width() )
				setPixelGreyBright( _image, i, j, 255 );
		}
	}
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::findLinesFrames( float _lineExpFactor, float _lineTopFactor, float _lineBotFactor )
{	
	FloatV rowsBrightness = getImageRowsBrightness( m_image );

	float averageBrightness = getAverageBrightness( rowsBrightness );
	float topBrightness = _lineTopFactor * averageBrightness;
	float botBrightness = _lineBotFactor * averageBrightness;	
	
	QPoint startPoint, endPoint;	
	bool startDetected = false;

	for ( size_t i = 0; i < rowsBrightness.size(); i++ )
	{
		if ( !startDetected && isTopRow( i, rowsBrightness, topBrightness, botBrightness ) )
		{	
			startDetected = true;			
			startPoint = QPoint( 1, i );			
		}
		if ( startDetected && isBotRow( i, rowsBrightness, topBrightness, botBrightness ) )
		{
			startDetected = false;
			endPoint = QPoint( m_image.width() - 2, i );
			m_linesFrames.push_back( QRect( startPoint, endPoint ) );			
			IntV layout;
			m_symbolsLayout.push_back( layout );
		}
	}	
	expandLinesBorders( _lineExpFactor );
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isTopRow( int _currRow, FloatV& _rowsBrightness, float _topBright, float _botBright )
{
	if ( ( _currRow - 2 < 0 ) || ( (unsigned)( _currRow + 3 ) >= _rowsBrightness.size() ) )
		return false;
	return (( _rowsBrightness[ _currRow - 2 ] < _topBright )
		&&	( _rowsBrightness[ _currRow - 1 ] < _topBright )
		&&	( _rowsBrightness[ _currRow + 0 ] > _topBright )
		&&	( _rowsBrightness[ _currRow + 1 ] > _botBright )
		&&	( _rowsBrightness[ _currRow + 2 ] > _botBright )
		&&	( _rowsBrightness[ _currRow + 3 ] > _botBright )
		);	
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isBotRow( int _currRow, FloatV& _rowsBrightness, float _topBright, float _botBright )
{
	if ( (unsigned)( _currRow + 3 ) >= _rowsBrightness.size() )
		return false;
	return (	(	( _rowsBrightness[ _currRow + 0 ] > _topBright )
				&&	( _rowsBrightness[ _currRow + 1 ] < _botBright )
				)
				||
				(
					( _rowsBrightness[ _currRow + 1 ] < _botBright )
				&&	( _rowsBrightness[ _currRow + 2 ] < _botBright )
				&&	( _rowsBrightness[ _currRow + 3 ] < _botBright )
				)
	);
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::expandLinesBorders( float _lineExpFactor )
{
	/** find minimum height */
	int minHeight = m_linesFrames[0].height();
	for ( size_t i = 1; i < m_linesFrames.size(); i++ )
	{
		if ( m_linesFrames[i].height() < minHeight )
			minHeight = m_linesFrames[i].height();
	}

	/** expanding borders */	
	int expandVal = _lineExpFactor * minHeight;
	for ( size_t i = 0; i < m_linesFrames.size(); i++ )
	{
		if ( ( m_linesFrames[i].top() - expandVal ) <= 0 )
			m_linesFrames[i].setTop( 0 );
		else
			m_linesFrames[i].setTop( m_linesFrames[i].top() - expandVal );
		if ( ( m_linesFrames[i].bottom() + expandVal ) >= m_image.height() )
			m_linesFrames[i].setBottom( m_image.height() - 1 );
		else
			m_linesFrames[i].setBottom( m_linesFrames[i].bottom() + expandVal );
	}
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::generateLinesImages()
{
	QPainter painter;
	for ( size_t i = 0; i < m_linesFrames.size(); i++ )
	{		
		QImage currImg( m_linesFrames[i].width(), m_linesFrames[i].height(), m_image.format() );
		painter.begin( &currImg );		
		painter.setRenderHint( QPainter::TextAntialiasing, true );		
		painter.drawImage( QPoint(0, 0), m_image, m_linesFrames[i] );
		painter.end();
		m_linesImages.push_back( currImg );
	}

	/** to be removed */
	//for ( size_t i = 0; i < m_linesImages.size(); i++ )
	//{
	//	QString name;
	//	name.setNum( i );
	//	m_linesImages[i].save( "line_" + name + ".bmp", "BMP" );
	//}
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::findWordsFrames( int _contrast, float _wordLeftFactor, float _wordRightFactor )
{
	QImageV tempImages;

	/** increasing line images contrast */
	for ( size_t t = 0; t < m_linesImages.size(); t++ )		
		tempImages.push_back( contrastIncrease( m_linesImages[t], _contrast ) );			

	/** painting nearby pixels */	
	for ( size_t t = 0; t < tempImages.size(); t++ )	
		tempImages[t] = imageDegradation( tempImages[t] );				

	/** processing each processed line image */
	for ( size_t t = 0; t < tempImages.size(); t++ )	
	{
		FloatV lineColsBrightness = getImageColsBrightness( tempImages[t] );

		float averageBrightness = getAverageBrightness( lineColsBrightness );
		float leftBrightness  = _wordLeftFactor  * averageBrightness;
		float rightBrightness = _wordRightFactor * averageBrightness;
	
		QPoint leftPoint, rightPoint;
		bool startDetected = false;

		for ( size_t i = 0; i < lineColsBrightness.size(); i++ )
		{
			if ( !startDetected && isLeftCol( i, lineColsBrightness, leftBrightness ) )
			{
				startDetected = true;
				leftPoint = QPoint( i - 1, 0 );
			}
			if ( startDetected && isRightCol( i, lineColsBrightness, rightBrightness ) )
			{
				startDetected = false;
				rightPoint = QPoint( i + 1, tempImages[t].height() - 1 );
				QPoint lineStart( m_linesFrames[t].topLeft() );
				m_wordsFrames.push_back( QRect( lineStart + leftPoint, lineStart + rightPoint ) );
				int layout;
				m_symbolsLayout[t].push_back( layout );
			}
		}
	}	
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isLeftCol( int _currPos, FloatV& _wordsBright, float _leftBright )
{
	if ( ( _currPos - 1 < 0 ) || ( (unsigned)( _currPos + 1 ) >= _wordsBright.size() ) )
		return false;
	return (( _wordsBright[ _currPos - 1 ] < _leftBright )
		&&	( _wordsBright[ _currPos + 0 ] > _leftBright )		
		&&	( _wordsBright[ _currPos + 1 ] > _leftBright )		
		);	
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isRightCol( int _currPos, FloatV& _wordsBright, float _rightBright )
{
	//if ( ( _currPos - 2 < 0 ) || ( (unsigned)( _currPos + 4 ) >= _wordsBright.size() ) )
	if ( ( _currPos - 1 < 0 ) || ( (unsigned)( _currPos + 1 ) >= _wordsBright.size() ) )
		return false;
	return ( /*( _wordsBright[ _currPos - 2 ] > _leftBright )*/
			( _wordsBright[ _currPos - 1 ] > _rightBright )
		&&	( _wordsBright[ _currPos + 0 ] < _rightBright )
		&&	( _wordsBright[ _currPos + 1 ] < _rightBright )
		//&&	( _wordsBright[ _currPos + 2 ] < _rightBright )
		//&&	( _wordsBright[ _currPos + 3 ] < _rightBright )
		//&&	( _wordsBright[ _currPos + 4 ] < _rightBright )
		);
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::generateWordsImages()
{
	QPainter painter;
	QImage contrImg = contrastIncrease( m_image, 25 );
	for ( size_t i = 0; i < m_wordsFrames.size(); i++ )
	{		
		QImage currImg( m_wordsFrames[i].width(), m_wordsFrames[i].height(), m_image.format() );
		painter.begin( &currImg );		
		painter.setRenderHint( QPainter::TextAntialiasing, true );		
		painter.drawImage( QPoint(0, 0), contrImg, m_wordsFrames[i] );
		painter.end();
		m_wordsImages.push_back( currImg );
	}

	/** to be removed */
	//for ( size_t i = 0; i < 5; i++ )
	//{
	//	QString name;
	//	name.setNum( i );
	//	m_wordsImages[i].save( "word_" + name + ".bmp", "BMP" );
	//}
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::findSymbolsFrames( float _intervalFactor, float _brightFactor, float _widthFactor )
{	
	/** process each word image */
	int currWordNum = 0;
	for ( size_t l = 0; l < m_symbolsLayout.size(); l++ )
	{		
		for ( size_t w = 0; w < m_symbolsLayout[l].size(); w++ )
		{
			/** finding bright parameters */
			FloatV wordColsBrightness = getImageColsBrightness( m_wordsImages[currWordNum] );			
			float averageBrightness = getAverageBrightness( wordColsBrightness );			
			float symbolBrightBorder = _brightFactor * averageBrightness;			
			
			/** finding minimum interval width */
			int intervalWidth =	_intervalFactor * m_wordsImages[currWordNum].height();			

			/** finding all local minimums */
			IntV localMinimums = findLocalMinimums( wordColsBrightness, intervalWidth );

			/** remove false minimums */
			localMinimums = checkLocalMinimums( wordColsBrightness, localMinimums, symbolBrightBorder );

			/** remove close minimums */	
			float minWidth = _widthFactor * m_wordsImages[currWordNum].height();
			localMinimums = removeCloseMinimums( localMinimums, minWidth );

			/** storing frames */
			int symbCount = 0;
			QPoint wordStart( m_wordsFrames[currWordNum].topLeft() );
			for ( size_t i = 0; i < localMinimums.size() - 1; i++ )
			{
				QPoint startPoint( localMinimums[i], 0 );
				QPoint endPoint( localMinimums[i+1], m_wordsFrames[currWordNum].height() - 1 );
				m_symbolsFrames.push_back( QRect( wordStart + startPoint, wordStart + endPoint ) );
				symbCount++;
			}

			/** store symbols count in word */
			m_symbolsLayout[l][w] = symbCount;

			/** next image */
			currWordNum++;			
		}
	}
}


/*----------------------------------------------------------------------------*/


IntV ImageSegmentator::findLocalMinimums( FloatV& _wordColsBright, int _interval )
{
	IntV minimums;

	int intervalStart = 0;
	int intervalEnd   = intervalStart + _interval;
	
	minimums.push_back( intervalStart );
	while ( (unsigned)intervalEnd < _wordColsBright.size() )
	{
		int minimum = findMinimumBright( _wordColsBright, intervalStart, intervalEnd );		
		minimums.push_back( minimum );		

		/** set next interval */
		intervalStart = minimum + 1;
		intervalEnd   = intervalStart + _interval;
	}
	minimums.push_back( _wordColsBright.size() - 1 );
	return minimums;
}


/*----------------------------------------------------------------------------*/


int ImageSegmentator::findMinimumBright( FloatV& _wordColsBright, int _startPos, int _endPos )
{
	int currPos = _startPos;
	float minBright = _wordColsBright[currPos];

	for ( int i = currPos + 1; i <= _endPos; i++ )
	{
		if ( _wordColsBright[i] <= minBright )
		{
			minBright = _wordColsBright[i];
			currPos = i;
		}
	}
	return currPos;
}


/*----------------------------------------------------------------------------*/


IntV ImageSegmentator::checkLocalMinimums( FloatV& _wordColsBright, IntV& _minimums, float _bright )
{
	IntV newLocalMinimums;

	newLocalMinimums.push_back( 0 );
	for ( size_t i = 0; i < _minimums.size(); i++ )
	{
		if ( isMinimum( _minimums[i], _wordColsBright, _bright ) )
			newLocalMinimums.push_back( _minimums[i] );
	}	
	newLocalMinimums.push_back( _minimums.back() );
	return newLocalMinimums;
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isMinimum( int _currPos, FloatV& _wordColsBright, float _symBrightBorder )
{	
	return ( ( _wordColsBright[ _currPos + 0 ] < _symBrightBorder ) );
}


/*----------------------------------------------------------------------------*/


IntV ImageSegmentator::removeCloseMinimums( IntV& _localMinimums, int _minWidth )
{
	IntV newLocalMinimums;	
	
	int startPos = 0;
	int endPos = 1;	
	newLocalMinimums.push_back( _localMinimums[startPos] );
	while( (unsigned)(endPos) < _localMinimums.size() )
	{		
		int currWidth = _localMinimums[endPos] - _localMinimums[startPos];		
		if ( currWidth < _minWidth )
		{
			endPos++;
		}
		else
		{
			newLocalMinimums.push_back( _localMinimums[endPos] );
			startPos = endPos;
			endPos = startPos + 1;
		}
	}
	return newLocalMinimums;
}


/*----------------------------------------------------------------------------*/


void ImageSegmentator::generateSymbolsImages( int _imageSize, int _contrast )
{	
	QPainter painter;
	QImage contrImg( contrastIncrease( m_image, _contrast ) );
	//contrImg.save( "test.bmp", "BMP" );
	for ( size_t i = 0; i < m_symbolsFrames.size(); i++ )
	{	
		/** generate images from frames */
		QImage currImg( m_symbolsFrames[i].width(), m_symbolsFrames[i].height(), m_image.format() );
		painter.begin( &currImg );		
		painter.setRenderHint( QPainter::TextAntialiasing, true );		
		painter.drawImage( QPoint(0, 0), contrImg, m_symbolsFrames[i] );
		painter.end();

		/** resizing and scaling image */
		//QString num;
		//num.setNum( i );
		//currImg.save( "temp/" + num + ".bmp", "BMP" );
		QImage readyImage = prepareImage( currImg, _imageSize );		
		m_symbolsImages.push_back( readyImage );		
	}
}


/*----------------------------------------------------------------------------*/


QImage ImageSegmentator::prepareImage( QImage& _currImage, int _imageSize )
{
	QRect imageFrame( _currImage.rect() );		

	/** top and bottom detect */
	FloatV rowsBrightness = getImageRowsBrightness( _currImage );	
	float rowsAvrBrightness = getAverageBrightness( rowsBrightness );	
	float topBrightness = 0.3 * rowsAvrBrightness;
	float botBrightness = 0.2 * rowsAvrBrightness;		

	for ( size_t i = 0; i < rowsBrightness.size(); i++ )
	{
		if ( isStartBorder( i, rowsBrightness, topBrightness ) )		
		{									
			imageFrame.setTop( i );			
			break;
		}
	}
	for ( size_t i = rowsBrightness.size() - 1; i > (unsigned)imageFrame.top(); i-- )
	{
		if( isEndBorder( i, rowsBrightness, botBrightness ) )
		{
			imageFrame.setBottom( i - 1 );			
			break;
		}
	}

	/** left and right detect */
	FloatV colsBrightness = getImageColsBrightness( _currImage );
	float colsAvrBrightness = getAverageBrightness( colsBrightness );
	float leftBrightness  = 0.3 * colsAvrBrightness;
	float rightBrightness = 0.2 * colsAvrBrightness;

	for ( size_t i = 0; i < colsBrightness.size(); i++ )
	{
		if ( isStartBorder( i, colsBrightness, leftBrightness ) )
		{
			imageFrame.setLeft( i );			
			break;
		}
	}
	for ( size_t i = colsBrightness.size() - 1; i > (unsigned)imageFrame.left(); i-- )
	{
		if ( isEndBorder( i, colsBrightness, rightBrightness ) )
		{
			imageFrame.setRight( i - 1 );			
			break;
		}
	}

	/** draw part */
	QPainter painter;

	/** cut image */
	QImage tempImage( imageFrame.width(), imageFrame.height(), _currImage.format() );
	painter.begin( &tempImage );
	painter.setRenderHint( QPainter::TextAntialiasing, true );
	painter.eraseRect( tempImage.rect() );
	painter.drawImage( QPoint(0, 0), _currImage, imageFrame );
	painter.end();	

	/** resize image */
	QImage scaledImage = tempImage.scaled( _imageSize, _imageSize, Qt::KeepAspectRatio );
	QImage readyImage( _imageSize, _imageSize, tempImage.format() );
	painter.begin( &readyImage );
	painter.setRenderHint( QPainter::TextAntialiasing, true );		
	painter.eraseRect( readyImage.rect() );
	painter.drawImage( QPoint( 0, 0 ), scaledImage );
	painter.end();	

	return readyImage;
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isStartBorder( int _currPos, FloatV& _brightness, float _bright )
{
	if ( ( _currPos - 1 < 0 ) || ( (unsigned)( _currPos + 1 ) >= _brightness.size() ) )
		return false;
	return (	( _brightness[ _currPos - 1 ] < _bright )
			&&	( _brightness[ _currPos + 0 ] > _bright )		
			&&	( _brightness[ _currPos + 1 ] > _bright )
	);
}


/*----------------------------------------------------------------------------*/


bool ImageSegmentator::isEndBorder( int _currPos, FloatV& _brightness, float _bright )
{
	if ( ( _currPos - 1 < 0 ) || ( (unsigned)( _currPos + 1 ) >= _brightness.size() ) )
		return false;
	return (	( _brightness[ _currPos - 1 ] > _bright )
			&&	( _brightness[ _currPos - 0 ] < _bright )
			&&	( _brightness[ _currPos + 1 ] < _bright )
	);
}


/*----------------------------------------------------------------------------*/

